import os
import os.path as osp
import json
import cv2
import numpy as np

label_dict = {
    'A': 0.0,
    'B': 1.0,
    'C': 2.0,
    'D': 3.0,
    'E': 4.0,
    'F': 5.0,
    'G': 6.0,
    'H': 7.0,
    'I': 8.0,
    'J': 9.0,
    'K': 10.0,
}

srcDIR = osp.abspath('/diskb/TZB/TZB_1_Data/val_json')  # 源文件夹
desDIR = osp.join(srcDIR, '..', 'val_txt')  # 输出文件夹
if not osp.exists(desDIR):
    os.makedirs(desDIR)
jsonfiles = [name for name in os.listdir(srcDIR) if name.endswith('.json')]
# print(jsonfiles)
# print(len(jsonfiles))

for filename in jsonfiles:  # labelme格式，一个json对应一个图片，对应输出一个txt
    path_json = osp.join(srcDIR, filename)
    with open(path_json, 'r') as fp:
        jsonData = fp.read()
    info = json.loads(jsonData)
    # print(info)
    info_shapes = info['shapes']
    info_imgpath = info['imagePath']
    info_imgH = info['imageHeight']
    info_imgW = info['imageWidth']
    # print(info_imgpath)
    OUTTXT = osp.join(desDIR, info_imgpath[:-4] + '.txt')
    with open(OUTTXT, 'w') as fp:  # 清空old标签
        fp.write('')

    for i in range(len(info_shapes)):  # 对每一个边界框
        label_i = info_shapes[i]['label']
        points_i = info_shapes[i]['points']  # xy4
        # print(f"{label_i}\n{points_i}")

        label_i = label_dict[label_i]

        rect_i = cv2.minAreaRect(np.int0(points_i))
        # print(rect_i)  # 旋转矩形((x,y),(w,h),a)
        (x, y), (w, h), a = rect_i
        # print(x, y, w, h, a)
        x = x / info_imgW  # 归一化
        y = y / info_imgH
        w = w / info_imgW
        h = h / info_imgH
        a = int(a)
        if a < 0:  # 度数范围0-179
            a += 180
        if w < h:  # 规定w是长边，h是短边
            w, h = h, w
            a += 90
            if a >= 180:  # 度数范围0-179
                a -= 180

        # print(x, y, w, h, a)
        # rect_i = (x, y), (w, h), a

        # filtering...  滤除不合法的边界框
        if x < 0 or x > 1.0 or y < 0 or y > 1.0:
            continue

        item = label_i, x, y, w, h, a
        item_str = ' '.join([str(i) for i in item])
        # print(item_str)

        with open(OUTTXT, 'a') as fp:
            fp.write(item_str)
            fp.write('\n')

    print(info_imgpath[:-4], 'done.')
